import os
import re
project_path = r"D:/data/android/demo/aidc"
cryption_cpp_file = os.path.join(
    project_path, "nativeLibs/app/src/main/cpp/cryption/model_file_decryption_iv.cpp"
)
if os.path.exists(cryption_cpp_file):
    print("cryption_cpp_file exists")
else:
    print("cryption_cpp_file not exists")
def addCPPFile(cryption_cpp_file:str, model_name:str,iv:str):
    """扫描所有 IV：包括 constexpr 和 map 初始化"""
    if not os.path.exists(cryption_cpp_file):
        print(f"❌ 文件不存在: {cryption_cpp_file}")
        return
    with open(cryption_cpp_file, 'r', encoding='utf-8') as f:
        content = f.read()
    pass
def scan_all_ivs(file_path):
    """扫描所有 IV：包括 constexpr 和 map 初始化"""
    if not os.path.exists(file_path):
        print(f"❌ 文件不存在: {file_path}")
        return [], []

    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    # === 1. 扫描 constexpr const char* IV ===
    pattern_constexpr = re.compile(
        r'(static\s+constexpr\s+const\s+char\*\s+)(\w+_MODEL_FILE_IV)\s*=\s*"([0-9a-fA-F]{32})";',
        re.MULTILINE
    )
    simple_ivs = []
    for match in pattern_constexpr.finditer(content):
        var_name = match.group(2)
        iv_value = match.group(3)
        start, end = match.span()
        line_no = content[:start].count('\n') + 1
        simple_ivs.append({
            'type': 'constexpr',
            'name': var_name,
            'value': iv_value,
            'start': start,
            'end': end,
            'line': line_no,
            'match': match
        })

    # === 2. 扫描 map 初始化中的 IV ===
    pattern_map = re.compile(
        r'\{\s*AreaType::(\w+),\s*"([0-9a-fA-F]{32})"\s*\}',
        re.MULTILINE
    )
    map_ivs = []
    for match in pattern_map.finditer(content):
        enum_name = match.group(1)
        iv_value = match.group(2)
        start, end = match.span()
        line_no = content[:start].count('\n') + 1
        map_ivs.append({
            'type': 'map',
            'enum': enum_name,
            'value': iv_value,
            'start': start,
            'end': end,
            'line': line_no,
            'match': match
        })

    return simple_ivs, map_ivs

def display_ivs(simple_ivs, map_ivs):
    """打印所有找到的 IV"""
    print("\n🔍 找到的 constexpr IV 变量:")
    print("-" * 80)
    for idx, iv in enumerate(simple_ivs, 1):
        print(f"{idx:2d}. {iv['name']:50s} = \"{iv['value']}\"  (第 {iv['line']} 行)")

    print("\n🔍 找到的 map 初始化 IV:")
    print("-" * 80)
    for idx, iv in enumerate(map_ivs, 1):
        print(f"{idx:2d}. {{ AreaType::{iv['enum']}, \"{iv['value']}\" }}  (第 {iv['line']} 行)")

def select_ivs_to_update(simple_ivs, map_ivs):
    """让用户选择要更新的 IV"""
    print("\n📌 选择要更新的 IV 类型:")
    print("  1) constexpr 字符串 IV")
    print("  2) map 初始化 IV")
    print("  3) 全部")
    choice = input("👉 请输入选择 (1/2/3): ").strip()

    selected_simple = []
    selected_map = []

    if choice == '1':
        if not simple_ivs:
            print("⚠️ 无 constexpr IV 可选。")
            return [], []
        selected_simple = prompt_select(simple_ivs, "constexpr")
    elif choice == '2':
        if not map_ivs:
            print("⚠️ 无 map IV 可选。")
            return [], []
        selected_map = prompt_select(map_ivs, "map")
    elif choice == '3':
        if simple_ivs:
            selected_simple = prompt_select(simple_ivs, "constexpr")
        if map_ivs:
            selected_map = prompt_select(map_ivs, "map")
    else:
        print("❌ 无效输入。")
    return selected_simple, selected_map

def prompt_select(iv_list, label):
    """通用选择逻辑"""
    display_list = []
    for i, iv in enumerate(iv_list, 1):
        if iv['type'] == 'constexpr' or iv['type'] == 'const_val':
            display_list.append(f"{i:2d}. {iv['name']}")
        else:  # map 类型
            display_list.append(f"{i:2d}. {iv['enum']}")
    print(f"\n✅ 可用的 {label} IV:")
    print('\n'.join(display_list))
    sel = input(f"请输入要更新的编号 (如 1,3,5 或 all): ").strip()
    if sel == 'all':
        return list(range(len(iv_list)))
    try:
        indices = [int(x.strip()) - 1 for x in sel.split(',') if x.strip().isdigit()]
        return [i for i in indices if 0 <= i < len(iv_list)]
    except ValueError:
        print("⚠️ 输入格式错误。")
        return []

def update_ivs(file_path, output_path,new_iv, simple_ivs, map_ivs, sel_simple, sel_map):
    """更新选中的 IV 并写入新文件"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    changes = []

    # 更新 constexpr IV
    for idx in sel_simple:
        iv = simple_ivs[idx]
        old_text = f'"{iv["value"]}"'
        new_text = f'{new_iv}'
        start = iv['match'].start(3)  # 第3组是 IV 值
        end = iv['match'].end(3)
        content = content[:start] + new_text + content[end:]
        changes.append({
            'type': 'constexpr',
            'name': iv['name'],
            'old': iv['value'],
            'new': new_iv,
            'line': iv['line']
        })

    # 更新 map IV
    for idx in sel_map:
        iv = map_ivs[idx]
        old_text = f'"{iv["value"]}"'
        new_text = f'{new_iv}'
        start = iv['match'].start(2)
        end = iv['match'].end(2)
        content = content[:start] + new_text + content[end:]
        changes.append({
            'type': 'map',
            'enum': iv['enum'],
            'old': iv['value'],
            'new': new_iv,
            'line': iv['line']
        })

    # 写入文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(content)

    # 打印日志
    print(f"\n✅ 共更新 {len(changes)} 个 IV 值:")
    for ch in changes:
        if ch['type'] == 'constexpr':
            print(f"🔄 {ch['name']:50s} = \"{ch['old']}\" → \"{ch['new']}\"")
        else:
            print(f"🔄 {{ AreaType::{ch['enum']:<20s}, \"{ch['old']}\" }} → \"{ch['new']}\"")

    print(f"\n💾 已保存到: {output_path}")
    return changes

def update_IV(newIv:str):
    print("🔍 正在扫描所有 IV 变量...")
    simple_ivs, map_ivs = scan_all_ivs(cryption_cpp_file)

    if not simple_ivs and not map_ivs:
        print("❌ 未找到任何 IV 变量，请检查文件格式。")
        return
    
    display_ivs(simple_ivs, map_ivs)
    sel_simple, sel_map = select_ivs_to_update(simple_ivs, map_ivs)

    if not sel_simple and not sel_map:
        print("👋 未选择任何 IV，退出。")
        return

    update_ivs(cryption_cpp_file, cryption_cpp_file,newIv, simple_ivs, map_ivs, sel_simple, sel_map)
    print("\n🎉 更新完成！")
def scan_all_model_files(file_path):
    """扫描 Kotlin 文件中的所有模型文件路径"""
    if not os.path.exists(file_path):
        print(f"❌ 文件不存在: {file_path}")
        return [], []

    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    # === 1. 扫描 const val 常量 ===
    pattern_const = re.compile(
        r'(const\s+val\s+)(\w+)\s*=\s*"([^"]*)";?',
        re.MULTILINE
    )
    simple_files = []
    for match in pattern_const.finditer(content):
        var_name = match.group(2)
        file_path_value = match.group(3)
        start, end = match.span()
        line_no = content[:start].count('\n') + 1
        simple_files.append({
            'type': 'const_val',
            'name': var_name,
            'value': file_path_value,
            'start': start,
            'end': end,
            'line': line_no,
            'match': match
        })

    # === 2. 扫描 val Map 类型 ===
    pattern_map = re.compile(
        r'AreaType\.(\w+)\s+to\s+"([^"]*)"',
        re.MULTILINE
    )
    map_files = []
    for match in pattern_map.finditer(content):
        enum_name = match.group(1)
        file_path_value = match.group(2)
        start, end = match.span()
        line_no = content[:start].count('\n') + 1
        map_files.append({
            'type': 'map',
            'enum': enum_name,
            'value': file_path_value,
            'start': start,
            'end': end,
            'line': line_no,
            'match': match
        })

    return simple_files, map_files

def display_model_files(simple_files, map_files):
    """打印所有找到的模型文件"""
    print("\n🔍 找到的 const val 模型文件:")
    print("-" * 80)
    for idx, file in enumerate(simple_files, 1):
        print(f"{idx:2d}. {file['name']:50s} = \"{file['value']}\"  (第 {file['line']} 行)")

    print("\n🔍 找到的 Map 模型文件:")
    print("-" * 80)
    for idx, file in enumerate(map_files, 1):
        print(f"{idx:2d}. {{ AreaType::{file['enum']}, \"{file['value']}\" }}  (第 {file['line']} 行)")

def update_model_files(file_path, output_path, new_path, simple_files, map_files, sel_simple, sel_map):
    """更新选中的模型文件路径并写入新文件"""
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    changes = []

    # 更新 const val 文件路径
    for idx in sel_simple:
        file = simple_files[idx]
        start = file['match'].start(3)  # 第3组是文件路径值
        end = file['match'].end(3)
        content = content[:start] + new_path + content[end:]
        changes.append({
            'type': 'const_val',
            'name': file['name'],
            'old': file['value'],
            'new': new_path,
            'line': file['line']
        })

    # 更新 map 文件路径
    for idx in sel_map:
        file = map_files[idx]
        start = file['match'].start(2)  # 第2组是文件路径值
        end = file['match'].end(2)
        content = content[:start] + new_path + content[end:]
        changes.append({
            'type': 'map',
            'enum': file['enum'],
            'old': file['value'],
            'new': new_path,
            'line': file['line']
        })

    # 写入文件
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(content)

    # 打印日志
    print(f"\n✅ 共更新 {len(changes)} 个模型文件路径:")
    for ch in changes:
        if ch['type'] == 'const_val':
            print(f"🔄 {ch['name']:50s} = \"{ch['old']}\" → \"{ch['new']}\"")
        else:
            print(f"🔄 {{ AreaType::{ch['enum']:<20s}, \"{ch['old']}\" }} → \"{ch['new']}\"")

    print(f"\n💾 已保存到: {output_path}")
    return changes
def update_x1_model(x1_project_path: str, new_model_path: str):
    """更新 X1 项目中的模型文件路径"""
    # decoder\src\main\java\com\action\ai\x1\decoder\model\ModelFile.kt
    model_file_path = os.path.join(x1_project_path, "decoder/src/main/java/com/action/ai/x1/decoder/model/ModelFile.kt")
    print(f"🔍 扫描 ModelFile.kt 文件...{model_file_path}")
    
    if not os.path.exists(model_file_path):
        print(f"❌ ModelFile.kt 文件不存在: {model_file_path}")
        return

    simple_files, map_files = scan_all_model_files(model_file_path)
    if not simple_files and not map_files:
        print("❌ 未找到任何模型文件路径，请检查文件格式。")
        return

    display_model_files(simple_files, map_files)
    sel_simple, sel_map = select_model_files_to_update(simple_files, map_files)

    if not sel_simple and not sel_map:
        print("👋 未选择任何模型文件，退出。")
        return

    update_model_files(model_file_path, model_file_path, new_model_path, simple_files, map_files, sel_simple, sel_map)
    print("\n🎉 X1 模型文件更新完成！")

def select_model_files_to_update(simple_files, map_files):
    """让用户选择要更新的模型文件"""
    print("\n📌 选择要更新的模型文件类型:")
    print("  1) const val 常量模型文件")
    print("  2) Map 类型模型文件")
    print("  3) 全部")
    choice = input("👉 请输入选择 (1/2/3): ").strip()

    selected_simple = []
    selected_map = []

    if choice == '1':
        if not simple_files:
            print("⚠️ 无 const val 模型文件可选。")
            return [], []
        selected_simple = prompt_select(simple_files, "const_val")
    elif choice == '2':
        if not map_files:
            print("⚠️ 无 Map 模型文件可选。")
            return [], []
        selected_map = prompt_select(map_files, "map")
    elif choice == '3':
        if simple_files:
            selected_simple = prompt_select(simple_files, "const_val")
        if map_files:
            selected_map = prompt_select(map_files, "map")
    else:
        print("❌ 无效输入。")
    return selected_simple, selected_map
def scan_model_quant_params(file_path):
    """
    扫描 Kotlin 文件中的模型量化参数
    
    Args:
        file_path: Kotlin 文件路径
    
    Returns:
        tuple: (simple_models, map_models) 
               simple_models: 简单列表类型的模型参数
               map_models: Map类型的模型参数
    """
    if not os.path.exists(file_path):
        print(f"❌ 文件不存在: {file_path}")
        return [], []

    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    # === 1. 扫描简单列表类型的模型 ===
    # 匹配 val MODEL_NAME = listOf(...)
    pattern_simple = re.compile(
        r'val\s+(\w+)\s*=\s*listOf\(\s*((?:mapOf\([^)]+\),?\s*)+)\)',
        re.MULTILINE | re.DOTALL
    )
    
    simple_models = []
    for match in pattern_simple.finditer(content):
        model_name = match.group(1)
        model_content = match.group(2)
        start, end = match.span()
        line_no = content[:start].count('\n') + 1
        
        # 提取参数对
        params = extract_params_from_content(model_content)
        
        simple_models.append({
            'type': 'simple',
            'name': model_name,
            'params': params,
            'start': start,
            'end': end,
            'line': line_no,
            'match': match
        })

    # === 2. 扫描 Map 类型的模型 ===
    # 使用更精确的正则表达式匹配 Map 类型模型
    # 使用更宽松的匹配方式
    pattern_map = re.compile(
        r'(val\s+(\w+):\s*Map<[^>]+>\s*=\s*mapOf\()((?:.|\n)*?\)\s*\))',
        re.MULTILINE | re.DOTALL
    )
    
    map_models = []
    for match in pattern_map.finditer(content):
        model_name = match.group(1)
        model_content = match.group(2)
        start, end = match.span()
        line_no = content[:start].count('\n') + 1
        
        # 提取各AreaType的参数
        area_params = extract_area_params_from_content(model_content)
        
        map_models.append({
            'type': 'map',
            'name': model_name,
            'params': area_params,
            'start': start,
            'end': end,
            'line': line_no,
            'match': match
        })

    return simple_models, map_models

def extract_params_from_content(content):
    """
    从模型内容中提取参数对
    
    Args:
        content: 模型内容字符串
    
    Returns:
        list: 参数对列表
    """
    params = []
    # 匹配 mapOf("key" to Pair(value, zero_point))
    param_pattern = re.compile(r'mapOf\("([^"]+)"\s+to\s+Pair\(([^,]+)f?,\s*(\d+)\)\s*\)')
    
    for match in param_pattern.finditer(content):
        key = match.group(1)
        scale = match.group(2)
        zero_point = match.group(3)
        
        params.append({
            'key': key,
            'scale': scale,
            'zero_point': zero_point
        })
    
    return params

def extract_area_params_from_content(content):
    """
    从Map类型模型内容中提取各AreaType的参数
    
    Args:
        content: 模型内容字符串
    
    Returns:
        dict: 各AreaType的参数字典
    """
    area_params = {}
    # 匹配 AreaType.TYPE to listOf(...)
    area_pattern = re.compile(r'AreaType\.(\w+)\s+to\s+listOf\(\s*((?:mapOf\([^)]+\)(?:,\s*)?)+)\s*\)')
    
    for match in area_pattern.finditer(content):
        area_type = match.group(1)
        area_content = match.group(2)
        
        # 提取该AreaType的参数
        params = extract_params_from_content(area_content)
        area_params[area_type] = params
    
    return area_params

def display_model_quant_params(simple_models, map_models):
    """
    显示模型量化参数
    
    Args:
        simple_models: 简单列表类型的模型参数
        map_models: Map类型的模型参数
    """
    print("\n🔍 找到的简单列表类型模型:")
    print("-" * 80)
    for idx, model in enumerate(simple_models, 1):
        print(f"{idx:2d}. {model['name']}  (第 {model['line']} 行)")
        for param in model['params']:
            print(f"     {param['key']}: scale={param['scale']}, zero_point={param['zero_point']}")

    print("\n🔍 找到的 Map 类型模型:")
    print("-" * 80)
    for idx, model in enumerate(map_models, 1):
        print(f"{idx:2d}. {model['name']}  (第 {model['line']} 行)")
        for area_type, params in model['params'].items():
            print(f"     AreaType.{area_type}:")
            for param in params:
                print(f"       {param['key']}: scale={param['scale']}, zero_point={param['zero_point']}")

# simple_ivs, map_ivs = scan_model_quant_params(r"D:\data\android\demo\aidc\x1ScannerApp\decoder\src\main\java\com\action\ai\x1\decoder\model\ModelQuantParam.kt")
# display_model_quant_params(simple_ivs, map_ivs)